import pandas as pd
import matplotlib.pyplot as plt

test_tsv = "/home/shaonian/SED/SED/configs/dataset/desed_tsv/val_strong_real.tsv"
test2_tsv= "/home/shaonian/SED/SED/configs/dataset/desed_tsv/test_strong.tsv"

test_df = pd.read_csv(test_tsv, sep="\t")
test_label = test_df["event_label"]
# all_test_labels = []
# for i in test_label:
#     for x in i.split(","):
#         all_test_labels.append(x)
# k_name = set(all_test_labels)
# test_label_count = {}
# for k in k_name:
#     test_label_count[k] = sum([x == k for x in all_test_labels])
test_label_count = test_label.value_counts()
test_label_count = test_label_count.to_dict()
test_label_count = {k.lower(): v for k, v in test_label_count.items()}

test2_df = pd.read_csv(test2_tsv, sep="\t")
test2_label = test2_df["event_label"]
test2_label_count = test2_label.value_counts()
test2_label_count = test2_label_count.to_dict()
test2_label_count = {k.lower(): v for k, v in test2_label_count.items()}

# plot the weak label counts
plt.figure(figsize=(10, 6))
# Convert weak_label_counts to percentages
test2_total = sum(test2_label_count.values())
test2_label_percentages = {k: (v/test2_total)*100 for k, v in test2_label_count.items()}

# Convert as_labels to percentages
test_total = sum(test_label_count.values())
test_label_percentages = {k: (v/test_total)*100 for k, v in test_label_count.items()}

all_labels = sorted(list(test_label_count.keys()))

# Prepare data for bar plot
test2_values = [test2_label_percentages.get(label, 0) for label in all_labels]
test_values = [test_label_percentages.get(label, 0) for label in all_labels]

x = range(len(all_labels))
width = 0.35

plt.bar([i - width/2 for i in x], test2_values, width, label='TestStrong', alpha=0.7)
plt.bar([i + width/2 for i in x], test_values, width, label='Selected', alpha=0.7)
plt.xticks(x, all_labels)
plt.title("Number of labels in non-single speech labels")
plt.xlabel("Labels")
plt.ylabel("Percentage (%)") 
plt.xticks(rotation=90)
plt.tight_layout()
plt.legend()
plt.savefig("../src/temp_results.png")